import cplex
import random
import numpy as np
import time
from pylab import *
from stoch_trnsport_gen import *
from cplex.exceptions import CplexError
import mpl_toolkits.mplot3d.axes3d as p3

NI = 20
NJ = 20
NW = 20
I = range(NI)
J = range(NJ)
O = range(NW)

prob,demand = gen_stoch_trnsport(NI, NJ, NW, eps=0.5)
prob.solve()
bigm = 0.1
true_opt_val = prob.solution.get_objective_value()

print prob.solution.get_objective_value(),
print prob.solution.progress.get_num_iterations()

X_bigm_vec = [0.02 * i for i in range(0,26)]
Y_bigm_vec = [0.02 * i for i in range(0,26)]
X_bigm_arr = np.array(X_bigm_vec)
Y_bigm_arr = np.array(Y_bigm_vec)
vals_arr = np.zeros((len(X_bigm_vec), len(X_bigm_vec)))
time_arr = np.zeros((len(X_bigm_vec), len(X_bigm_vec)))
x_bigm_star = 0
y_bigm_star = 0
best_use_time = 100000.0

Z = vals_arr
T = time_arr

#Iteratve over X_bigms and Y_bigms
for l in range(len(X_bigm_vec)):
    for k in range(len(Y_bigm_vec)):
        
        for w in O:
            for j in J:
                row_name = 'scenario' + str(w) + '_' + 'customer' + str(j)
                col_name = 'y' + '_' + str(w)
                
                if w%2 == 0:
                    bigm = X_bigm_vec[l]
                else:
                    bigm = Y_bigm_vec[k]

                prob.linear_constraints.set_coefficients(row_name, col_name, -bigm * demand[w][j])
                prob.linear_constraints.set_rhs(row_name, demand[w][j] * (1 - bigm) )

        start_time = time.time()
        prob.solve()
        use_time = time.time() - start_time
        #       vals_arr[l] = prob.solution.get_objective_value()
        time_arr[l, k] = use_time
        vals_arr[l, k] = prob.solution.get_objective_value()

        temp_opt_val = prob.solution.get_objective_value()

        if abs(temp_opt_val - true_opt_val) <= 0.000001 and use_time <= best_use_time:

            best_use_time = use_time
            x_bigm_star = X_bigm_vec[l]
            y_bigm_star = Y_bigm_vec[k]
            
        print prob.solution.get_objective_value(), 
        print prob.solution.progress.get_num_iterations(),
        print X_bigm_vec[l], Y_bigm_vec[k]

print "================="
print x_bigm_star, y_bigm_star
print best_use_time

X,Y = meshgrid(X_bigm_arr, Y_bigm_arr)
print X
print Y
print Z
print T

fig=figure()
ax = p3.Axes3D(fig)
ax.plot_wireframe(X,Y,Z)

#ax.plot_surface(X,Y,Z)
#ax.contourf3D(X, Y, Z)
#ax.contour3D(X, Y, Z)
#ax.plot(ravel(X),ravel(Y),ravel(Z))

ax.set_xlabel('Normalized big-M(1)')
ax.set_ylabel('Normalized big-M(2)')
ax.set_zlabel('Optimal cost')
show()


